#include "Sxai.h"
#include <cuda_runtime_api.h>
#include <cuda_fp16.h>
#include <stdio.h>
#include <cassert>
#include <fstream>
#include <iostream>
#include <thread>
#include "gemv.h"


namespace nvinfer1 {
namespace plugin {

const char *Sxai_PLUGIN_VERSION{"1"};
const char *Sxai_PLUGIN_NAME{"Sxai"};

Sxai::Sxai() { }

Sxai::Sxai(const void *data, size_t length) {
}

Sxai::~Sxai() { terminate(); }

int Sxai::getNbOutputs() const noexcept { return 1; }

DimsExprs Sxai::getOutputDimensions(int index, const DimsExprs *inputs, int nbInputDims,
                            nvinfer1::IExprBuilder &exprBuilder) noexcept {
    const auto & d0 = inputs[0].d[0]->getConstantValue();
    const auto & d1 = inputs[0].d[1]->getConstantValue();
    const auto & d2 = inputs[0].d[2]->getConstantValue();
    return nvinfer1::DimsExprs{5,{exprBuilder.constant(d0),exprBuilder.constant(d1), exprBuilder.constant(d2), exprBuilder.constant(1), exprBuilder.constant(1)}};
}

DataType Sxai::getOutputDataType(int index,
                                   const nvinfer1::DataType *inputTypes,
                                   int nbInputs) const noexcept {
    return mDataType;
}

int Sxai::initialize() noexcept { return STATUS_SUCCESS; }

void Sxai::terminate() noexcept {}

size_t Sxai::getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
                                int nbInputs,
                                const nvinfer1::PluginTensorDesc *outputs,
                                int nbOutputs) const noexcept {
  return 0;
}

int Sxai::enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
                    const nvinfer1::PluginTensorDesc *outputDesc,
                    const void *const *inputs, void *const *outputs,
                    void *workspace, cudaStream_t stream) noexcept {

    int d0 = inputDesc[0].dims.d[0];//600
    int d1 = inputDesc[0].dims.d[1];//200
    int d2 = inputDesc[0].dims.d[2];//8
    int d3 = inputDesc[0].dims.d[3];//1
    int d4 = inputDesc[0].dims.d[4];//64

    //printf("d0:%d, d1:%d, d2:%d, d3:%d, d4:%d\n", d0, d1, d2, d3, d4);

    //printf("inside Sxai enqueue...\n");
    if(mDataType == nvinfer1::DataType::kFLOAT){

        fgemv((float*)inputs[0], (float*)inputs[1], (float*)outputs[0], d0, d1, d4, stream);
    }else if(mDataType == nvinfer1::DataType::kHALF){

        //hgemv((half*)inputs[0], (half*)inputs[1], (half*)outputs[0], d0, d1, d4, stream);
    }

    return 0;
}

size_t Sxai::getSerializationSize() const noexcept {
  return 0;
}

void Sxai::serialize(void *buffer) const noexcept {

    return;
}

bool Sxai::supportsFormatCombination(int pos,
                                       const nvinfer1::PluginTensorDesc *inOut,
                                       int nbInputs, int nbOutputs) noexcept {
  ASSERT(inOut && pos < (nbInputs + nbOutputs));

  return (inOut[pos].type == inOut[0].type) && (inOut[pos].format == TensorFormat::kLINEAR);
}

const char *Sxai::getPluginType() const noexcept {
  return Sxai_PLUGIN_NAME;
}

const char *Sxai::getPluginVersion() const noexcept {
  return Sxai_PLUGIN_VERSION;
}

void Sxai::destroy() noexcept {delete this;}

nvinfer1::IPluginV2DynamicExt *Sxai::clone() const noexcept {
  auto plugin = new Sxai{};
  plugin->setPluginNamespace(namespace_);
  return plugin;
}

void Sxai::setPluginNamespace(const char *pluginNamespace) noexcept {
  namespace_ = pluginNamespace;
}

const char *Sxai::getPluginNamespace() const noexcept { return namespace_; }

void Sxai::attachToContext(cudnnContext *cudnn, cublasContext *cublas,
                             nvinfer1::IGpuAllocator *allocator) noexcept {

  mCublasHandle = cublas;
      
}

void Sxai::detachFromContext() noexcept {}

void Sxai::configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
                             int nbInputs,
                             const nvinfer1::DynamicPluginTensorDesc *out,
                             int nbOutputs) noexcept {
  auto type = in->desc.type;
  auto format = in->desc.format;
  assert(nbOutputs == 1);
  assert(nbInputs == 2);
  assert(type == DataType::kFLOAT || type == DataType::kHALF);
  mDataType = type;
}

PluginFieldCollection SxaiPluginCreator::field_collection_;
std::vector<PluginField> SxaiPluginCreator::attributes_;

SxaiPluginCreator::SxaiPluginCreator() {

}

const char *SxaiPluginCreator::getPluginName() const noexcept {
  return Sxai_PLUGIN_NAME;
}

const char *SxaiPluginCreator::getPluginVersion() const noexcept {
  return Sxai_PLUGIN_VERSION;
}

const PluginFieldCollection *SxaiPluginCreator::getFieldNames() noexcept {
  return &field_collection_;
}

IPluginV2DynamicExt *
SxaiPluginCreator::createPlugin(const char *name,
                                  const PluginFieldCollection *fc) noexcept {
  auto *plugin = new Sxai();
  plugin->setPluginNamespace(mNamespace.c_str());
  return plugin;
}

IPluginV2DynamicExt *
SxaiPluginCreator::deserializePlugin(const char *name, const void *serialData,
                                       size_t serialLength) noexcept {

  IPluginV2DynamicExt *plugin = new Sxai(serialData, serialLength);
  plugin->setPluginNamespace(mNamespace.c_str());
  return plugin;
}

} // namespace plugin
} // namespace nvinfer1